#!/usr/bin/env python3
"""
AI数据服务器主应用
提供多种登录方式和数据加密的API服务
"""

from flask import Flask, jsonify, request, g
from flask_jwt_extended import JWTManager, get_jwt
from flask_cors import CORS
import redis
from datetime import datetime
import json

from config import Config
from models import db
from routes import auth_bp, user_bp, wechat_bp
from routes.security import security_bp
from utils import EncryptionUtil
from security_middleware import SecurityMiddleware, get_security_headers
from security_enhanced import EnhancedSecurityMiddleware
from security_monitor import SecurityMonitor

# 全局加密工具实例
encryption_util = None

def create_app():
    """创建Flask应用"""
    app = Flask(__name__)
    app.config.from_object(Config)
    
    # 初始化全局加密工具
    global encryption_util
    encryption_util = EncryptionUtil(app.config['ENCRYPTION_KEY'])
    
    # 初始化Redis客户端
    redis_client = redis.from_url(app.config['REDIS_URL'])
    
    # 初始化安全中间件
    security_middleware = SecurityMiddleware(redis_client)
    enhanced_security_middleware = EnhancedSecurityMiddleware(redis_client)
    security_monitor = SecurityMonitor(redis_client)
    
    # 启动安全监控
    security_monitor.start_monitoring()
    
    # 初始化扩展
    db.init_app(app)
    
    # 初始化JWT
    jwt = JWTManager(app)
    
    # 初始化CORS
    CORS(app, origins=app.config['CORS_ORIGINS'])
    
    # 全局安全中间件 - 在请求处理前设置安全上下文
    @app.before_request
    def setup_security_context():
        """设置安全上下文"""
        g.security_middleware = security_middleware
        g.enhanced_security_middleware = enhanced_security_middleware
        g.security_monitor = security_monitor
        g.redis_client = redis_client
    
    # 全局请求预处理中间件 - 解密请求数据
    @app.before_request
    def decrypt_request():
        """解密请求数据"""
        # 跳过健康检查和根路径
        if request.endpoint in ['health_check', 'index']:
            return
        
        # 只处理JSON请求
        if request.is_json and request.method in ['POST', 'PUT', 'PATCH']:
            try:
                data = request.get_json()
                
                # 检查是否是加密请求
                if data and data.get('encrypted', False):
                    encrypted_data = data.get('data')
                    if encrypted_data:
                        # 解密数据
                        decrypted_data = encryption_util.decrypt_json(encrypted_data)
                        
                        # 将解密后的数据存储到g对象中，供路由使用
                        g.decrypted_data = decrypted_data
                        g.is_encrypted_request = True
                    else:
                        g.decrypted_data = {}
                        g.is_encrypted_request = True
                else:
                    # 非加密请求，直接使用原始数据
                    g.decrypted_data = data or {}
                    g.is_encrypted_request = False
                    
            except Exception as e:
                return jsonify({
                    'success': False,
                    'message': f'请求数据解密失败: {str(e)}',
                    'timestamp': datetime.utcnow().isoformat()
                }), 400
        else:
            g.decrypted_data = {}
            g.is_encrypted_request = False
    
    # 全局响应后处理中间件 - 加密响应数据和添加安全头
    @app.after_request
    def encrypt_response(response):
        """加密响应数据并添加安全头"""
        # 添加安全响应头
        security_headers = get_security_headers()
        for header, value in security_headers.items():
            response.headers[header] = value
        
        # 跳过健康检查和根路径
        if request.endpoint in ['health_check', 'index']:
            return response
        
        # 只处理JSON响应
        if response.content_type and 'application/json' in response.content_type:
            try:
                # 检查是否需要加密响应（基于请求是否加密）
                if hasattr(g, 'is_encrypted_request') and g.is_encrypted_request:
                    response_data = response.get_json()
                    
                    if response_data:
                        # 加密响应数据
                        encrypted_data = encryption_util.encrypt_json(response_data)
                        
                        # 构建加密响应
                        encrypted_response = {
                            'encrypted': True,
                            'data': encrypted_data
                        }
                        
                        response.data = json.dumps(encrypted_response)
                        response.headers['Content-Length'] = len(response.data)
                        
            except Exception as e:
                # 如果加密失败，记录错误但不影响响应
                print(f"响应加密失败: {str(e)}")
        
        return response
    
    # JWT令牌黑名单检查
    @jwt.token_in_blocklist_loader
    def check_if_token_revoked(jwt_header, jwt_payload):
        jti = jwt_payload['jti']
        token_in_redis = redis_client.get(f"blacklist:{jti}")
        return token_in_redis is not None
    
    # JWT错误处理
    @jwt.expired_token_loader
    def expired_token_callback(jwt_header, jwt_payload):
        return jsonify({
            'success': False,
            'message': '令牌已过期',
            'timestamp': datetime.utcnow().isoformat()
        }), 401
    
    @jwt.invalid_token_loader
    def invalid_token_callback(error):
        return jsonify({
            'success': False,
            'message': '无效的令牌',
            'timestamp': datetime.utcnow().isoformat()
        }), 401
    
    @jwt.unauthorized_loader
    def missing_token_callback(error):
        return jsonify({
            'success': False,
            'message': '缺少访问令牌',
            'timestamp': datetime.utcnow().isoformat()
        }), 401
    
    @jwt.revoked_token_loader
    def revoked_token_callback(jwt_header, jwt_payload):
        return jsonify({
            'success': False,
            'message': '令牌已被撤销',
            'timestamp': datetime.utcnow().isoformat()
        }), 401
    
    # 注册蓝图
    app.register_blueprint(auth_bp, url_prefix='/api/auth')
    app.register_blueprint(user_bp, url_prefix='/api/user')
    app.register_blueprint(wechat_bp, url_prefix='/api/wechat')
    app.register_blueprint(security_bp, url_prefix='/api/security')
    
    # 健康检查接口
    @app.route('/health')
    def health_check():
        return jsonify({
            'status': 'healthy',
            'timestamp': datetime.utcnow().isoformat(),
            'version': '1.0.0'
        })
    
    # 根路径
    @app.route('/')
    def index():
        return jsonify({
            'message': 'AI数据服务器API',
            'version': '1.0.0',
            'timestamp': datetime.utcnow().isoformat(),
            'endpoints': {
                'health': '/health',
                'auth': '/api/auth',
                'user': '/api/user',
                'wechat': '/api/wechat',
                'security': '/api/security'
            }
        })
    
    # 全局错误处理
    @app.errorhandler(404)
    def not_found(error):
        return jsonify({
            'success': False,
            'message': '接口不存在',
            'timestamp': datetime.utcnow().isoformat()
        }), 404
    
    @app.errorhandler(500)
    def internal_error(error):
        return jsonify({
            'success': False,
            'message': '服务器内部错误',
            'timestamp': datetime.utcnow().isoformat()
        }), 500
    
    @app.errorhandler(400)
    def bad_request(error):
        return jsonify({
            'success': False,
            'message': '请求参数错误',
            'timestamp': datetime.utcnow().isoformat()
        }), 400
    
    # 请求日志中间件
    @app.before_request
    def log_request_info():
        if app.config['DEBUG']:
            print(f"[{datetime.utcnow()}] {request.method} {request.url}")
    
    return app

def get_request_data():
    """获取解密后的请求数据"""
    return getattr(g, 'decrypted_data', {})

def main():
    """主函数"""
    app = create_app()
    
    print("🚀 AI数据服务器启动中...")
    print(f"📍 服务地址: http://localhost:8000")
    print(f"🔧 调试模式: {'开启' if app.config['DEBUG'] else '关闭'}")
    print(f"💾 数据库: {app.config['SQLALCHEMY_DATABASE_URI']}")
    print(f"🔐 全局加解密: 已启用")
    print("📚 API文档:")
    print("   - 健康检查: GET /health")
    print("   - 认证接口: /api/auth/*")
    print("   - 用户接口: /api/user/*")
    print("   - 微信接口: /api/wechat/*")
    print("   - 安全管理: /api/security/*")
    print("=" * 50)
    
    # 启动应用
    app.run(
        host='0.0.0.0',
        port=8000,
        debug=app.config['DEBUG']
    )

if __name__ == '__main__':
    main()